#! /usr/bin/python

""" This script will determine the LFP_weighted_average stimulus of a set of experiments supplied to the command line"""

import os
import sys
import cPickle

from projects.electrophysiology.view_data.experiment import Experiment
from LFP import LFP_weighted_average as lwa

# DEFAULTS
tau = 0.300 # in sec
target_frequency = 40 # Hz
tr = 1.0/120.0 # in sec
fr = 1.0 # in Hz
sfs = 100 # in pixels
sff = 1.1 # factor (unitless)


# ---------------------- PARSE COMMAND LINE ARGUMENTS -----------
bath_exp_num_list=[]
exp_num_list=[]
for arg in sys.argv[1:]:
    if 'bath=' in arg:
        bath_exp_num = int(arg.split('=')[1])
        bath_exp_num_list.append(bath_exp_num)
    elif 'tau=' in arg:
        tau = float(arg.split('=')[1])
    elif 'target_frequency=' in arg:
        target_frequency = float(arg.split('=')[1])
    elif 'time_resolution=' in arg:
        tr = float(arg.split('=')[1])
    elif 'save_frames_factor=' in arg:
        sff = float(arg.split('=')[1])
    elif 'save_frames_size=' in arg:
        sfs = float(arg.split('=')[1])
    elif 'frequency_resolution=' in arg:
        fr = float(arg.split('=')[1])
    else:
        exp_num = int(arg)
        exp_num_list.append(exp_num)

print "\nExp_num_list is %s" % exp_num_list
print "The bath_exp_num_list is %s" % bath_exp_num_list
print "Tau = %f (seconds)" % tau
print "Time Resolution = %f (seconds)" % tr
print "Frequency Resolution = %f (Hz) " % fr
print "Save Frames Factor = %f (unitless)" % sff
print "Save Frames Size = %f (pixels)" % sfs
print "Target Frequency = %f (Hz)\n" % target_frequency


print "Opening experiments from the database"
exp_list = [Experiment.from_default_db(id) for id in exp_num_list]
bath_exp_list = [Experiment.from_default_db(id) for id in bath_exp_num_list]
print " -- Success."

# ----------------- CHECK INPUTS ---------------------------
print "Checking to make sure the experiments are all of the same stimulus"
exp_stim_type_list = [e.stim_type for e in exp_list]
exp_stim_type_set = set(exp_stim_type_list)
if len(exp_stim_type_set) == 1:
    print " --Passed."
else:
    print " --Failed, check your experiment inputs."
    sys.exit(1)

print "Checking to make sure the experiments are all of the same duration"
voltage_traces = [e.data[0] for e in exp_list]
vt_lengths = [len(vt) for vt in voltage_traces]
length_set = set(vt_lengths)
print "    experiment data lengths: %s" % vt_lengths
if len(length_set) == 1:
    print " --Passed."
else:
    print " --Failed, check your experiment inputs."
    sys.exit(1)

if len(bath_exp_list) > 1:
    print "Checking to make sure the bath data are all the same length."
    bath_data_list = [e.data for e in bath_exp_list]
    length_list = [len(bdl[0]) for bdl in bath_data_list]
    print "    bath_data lengths: %s" % length_list
    length_set = set(length_list)
    if len(length_set) == 1:
        print " --Passed."
    else:
        print " --Failed, check your bath inputs."
        sys.exit(1)

# ------------------ RUN LFP AVERAGE CALCULATION -------------
voltage_traces = [e.data[0] for e in exp_list]
first_experiment = exp_list[0]
sample_rate = first_experiment.data[1]
pulse_index = first_experiment.data[2]
frames = first_experiment.get_frames(sff, sfs)
frame_times_list = [e.frame_times for e in exp_list]
bath_traces = [e.data[0] for e in bath_exp_list]
print "Calculating LFP Weighted Average ..."
results, power = lwa.LFP_weighted_avg(voltage_traces, 
                                      sample_rate, 
                                      pulse_index,
                                      frames,
                                      frame_times_list,
                                      target_frequency,
                                      tau,
                                      time_resolution=tr,
                                      frequency_resolution=fr,
                                      noise_signals=bath_traces)

results_filename = "" 
for exp_num in sorted(exp_num_list):
    results_filename += '%d_' % exp_num 
results_filename += 'LFP_weighted_average.cPickle'

print "Saving Results to %s" % results_filename
result_dict = {'results':results, 'power':power}
result_dict['tau'] = tau
result_dict['target_frequency'] = target_frequency
result_dict['time_resolution'] = tr
result_dict['frequency_resolution'] = fr
result_dict['save_frames_size'] = sfs
result_dict['save_frames_factor'] = sff
result_dict['exp_num_list'] = exp_num_list
result_dict['bath_exp_num_list'] = bath_exp_num_list

with open(results_filename, 'w') as ofile:
    cPickle.dump(result_dict, ofile, protocol=-1)

